import logging
import os
import datetime
import colorama
from config.config import LOG_DIR, LOG_FORMAT

class TestLogger:
    def __init__(self):
        colorama.init()
        
        if not os.path.exists(LOG_DIR):
            os.makedirs(LOG_DIR)
            
        self.setup_logger()

    def setup_logger(self):
        timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
        log_file = os.path.join(LOG_DIR, f'test_{timestamp}.log')

        class ColoredFormatter(logging.Formatter):
            def __init__(self, fmt=None, datefmt=None):
                super().__init__(fmt, datefmt)
                self.GREEN = colorama.Fore.GREEN
                self.RED = colorama.Fore.RED
                self.YELLOW = colorama.Fore.YELLOW
                self.RESET = colorama.Fore.RESET

            def format(self, record):
                if record.levelno == logging.ERROR:
                    record.msg = f"{self.RED}{record.msg}{self.RESET}"
                elif "成功" in str(record.msg) or "Pass" in str(record.msg):
                    record.msg = f"{self.GREEN}{record.msg}{self.RESET}"
                else:
                    record.msg = f"{self.YELLOW}{record.msg}{self.RESET}"
                return super().format(record)

        logger = logging.getLogger('test_logger')
        logger.setLevel(logging.INFO)
        logger.handlers.clear()

        file_handler = logging.FileHandler(log_file, encoding='utf-8')
        file_handler.setFormatter(logging.Formatter(LOG_FORMAT))

        console_handler = logging.StreamHandler()
        console_handler.setFormatter(ColoredFormatter(LOG_FORMAT))

        logger.addHandler(file_handler)
        logger.addHandler(console_handler)

        self.logger = logger 